package cmds

import (
	"flag"
	"fmt"
	"time"
)

type FlagSet struct {
	*flag.FlagSet
	err         error
	requireArgs []string
}

type errRequireFlag struct {
	name  string
	usage string
}

func (e *errRequireFlag) Error() string {
	return fmt.Sprintf("argument [%s] is requied: %s", e.name, e.usage)
}

func NewFlagSet(name string, errorHandling flag.ErrorHandling) *FlagSet {
	return &FlagSet{
		FlagSet: flag.NewFlagSet(name, errorHandling),
	}
}

func (f *FlagSet) checkRequire(name string, require ...bool) {
	if len(require) > 0 && require[0] {
		f.requireArgs = append(f.requireArgs, name)
	}
}

func (f *FlagSet) Bool(name string, value bool, usage string, require ...bool) *bool {
	f.checkRequire(name, require...)
	return f.FlagSet.Bool(name, value, usage)
}

func (f *FlagSet) BoolVar(p *bool, name string, value bool, usage string, require ...bool) {
	f.checkRequire(name, require...)
	f.FlagSet.BoolVar(p, name, value, usage)
}

func (f *FlagSet) String(name string, value string, usage string, require ...bool) *string {
	f.checkRequire(name, require...)
	return f.FlagSet.String(name, value, usage)
}

func (f *FlagSet) StringVar(p *string, name string, value string, usage string, require ...bool) {
	f.checkRequire(name, require...)
	f.FlagSet.StringVar(p, name, value, usage)
}

func (f *FlagSet) Int(name string, value int, usage string, require ...bool) *int {
	f.checkRequire(name, require...)
	return f.FlagSet.Int(name, value, usage)
}

func (f *FlagSet) IntVar(p *int, name string, value int, usage string, require ...bool) {
	f.checkRequire(name, require...)
	f.FlagSet.IntVar(p, name, value, usage)
}

func (f *FlagSet) Int64(name string, value int64, usage string, require ...bool) *int64 {
	f.checkRequire(name, require...)
	return f.FlagSet.Int64(name, value, usage)
}

func (f *FlagSet) Int64Var(p *int64, name string, value int64, usage string, require ...bool) {
	f.checkRequire(name, require...)
	f.FlagSet.Int64Var(p, name, value, usage)
}

func (f *FlagSet) Float64(name string, value float64, usage string, require ...bool) *float64 {
	f.checkRequire(name, require...)
	return f.FlagSet.Float64(name, value, usage)
}

func (f *FlagSet) Float64Var(p *float64, name string, value float64, usage string, require ...bool) {
	f.checkRequire(name, require...)
	f.FlagSet.Float64Var(p, name, value, usage)
}

func (f *FlagSet) Duration(name string, value time.Duration, usage string, require ...bool) *time.Duration {
	f.checkRequire(name, require...)
	return f.FlagSet.Duration(name, value, usage)
}

func (f *FlagSet) DurationVar(p *time.Duration, name string, value time.Duration, usage string, require ...bool) {
	f.checkRequire(name, require...)
	f.FlagSet.DurationVar(p, name, value, usage)
}

func (f *FlagSet) Func(name, usage string, fn func(string) error, require ...bool) {
	f.checkRequire(name, require...)
	f.FlagSet.Func(name, usage, fn)
}

func (f *FlagSet) UintVar(p *uint, name string, value uint, usage string, require ...bool) {
	f.checkRequire(name, require...)
	f.FlagSet.UintVar(p, name, value, usage)
}

func (f *FlagSet) Uint(name string, value uint, usage string, require ...bool) *uint {
	f.checkRequire(name, require...)
	return f.FlagSet.Uint(name, value, usage)
}

func (f *FlagSet) Uint64Var(p *uint64, name string, value uint64, usage string, require ...bool) {
	f.checkRequire(name, require...)
	f.FlagSet.Uint64Var(p, name, value, usage)
}

func (f *FlagSet) Uint64(name string, value uint64, usage string, require ...bool) *uint64 {
	f.checkRequire(name, require...)
	return f.FlagSet.Uint64(name, value, usage)
}

func (f *FlagSet) Usage() {
	if f.err != nil {
		fmt.Println("Error:", f.err.Error())
	}
	f.FlagSet.Usage()
}

func (f *FlagSet) Parse(arguments []string) error {
	if err := f.FlagSet.Parse(arguments); err != nil {
		return err
	}

	actual := map[string]*flag.Flag{}
	f.FlagSet.Visit(func(f *flag.Flag) {
		actual[f.Name] = f
	})

	for _, a := range f.requireArgs {
		if _, ok := actual[a]; !ok {
			flag := f.FlagSet.Lookup(a)
			f.err = &errRequireFlag{
				name:  flag.Name,
				usage: flag.Usage,
			}
			return f.err
		}
	}

	return f.FlagSet.Parse(arguments)
}
